#pragma once
#include <iostream>
#include <codecvt>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <onnxruntime_c_api.h>
#include <onnxruntime_cxx_api.h>
#include<vector>
#include<string>
using std::string;

struct Detection
{
    cv::Rect box;
    float conf;
    int classId;
    string className;
};

class Yolo5DetetorCpp
{
public:
    static void LogFuntion(void* param, OrtLoggingLevel severity, const char* category, const char* logid, const char* code_location, const char* message);
    explicit Yolo5DetetorCpp(std::nullptr_t) {};
    Yolo5DetetorCpp(const std::string& modelPath,
        const bool isGPU,
        const cv::Size& inputSize);

    std::vector<Detection> detect(cv::Mat& image, const float& confThreshold, const float& iouThreshold);
    std::vector<Detection> detect(string imagepath, const float& confThreshold, const float& iouThreshold);

    std::vector<string> classNames;
private:
    Ort::Env env{ nullptr };
    Ort::SessionOptions sessionOptions{ nullptr };
    Ort::Session session{ nullptr };

    void preprocessing(cv::Mat& image, float*& blob, std::vector<int64_t>& inputTensorShape);
    std::vector<Detection> postprocessing(const cv::Size& resizedImageShape,
        const cv::Size& originalImageShape,
        std::vector<Ort::Value>& outputTensors,
        const float& confThreshold, const float& iouThreshold);

    static void getBestClassInfo(std::vector<float>::iterator it, const int& numClasses,
        float& bestConf, int& bestClassId);

    std::vector<const char*> inputNames;
    std::vector<const char*> outputNames;
    bool isDynamicInputShape{};
    cv::Size2f inputImageShape;

};

